#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :ts文件抽帧.py
# @Time      :2024/11/7
# @Author    :CL
# @email     :1037654919@qq.com



import cv2
import os


def extract_frames(ts_file_path, output_dir, interval=5):
    # 使用 OpenCV 打开 TS 视频文件
    cap = cv2.VideoCapture(ts_file_path)

    # 检查视频是否成功打开
    if not cap.isOpened():
        print("无法打开视频文件。请检查路径是否正确。")
        return

    # 获取视频的帧率 (fps)
    fps = cap.get(cv2.CAP_PROP_FPS)
    fps = 30
    # # 获取视频的总帧数
    # total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # 获取视频时长 (单位：毫秒)，然后转为秒
    total_duration_ms = cap.get(cv2.CAP_PROP_POS_MSEC)  # 该值给出最后帧的时间戳（毫秒）
    total_duration_sec = total_duration_ms / 1000.0  # 转换为秒

    # 获取视频的宽度和高度
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # 获取视频的编码格式
    codec = cap.get(cv2.CAP_PROP_FOURCC)
    # # 将编码格式转换为字符（四个字节）
    # codec_str = ''.join([chr((codec >> (8 * i)) & 0xFF) for i in range(4)])

    # 打印视频文件的一般信息
    print(f"视频文件: {ts_file_path}")
    print(f"视频时长: {total_duration_sec:.2f} 秒")
    print(f"帧率: {fps:.2f} FPS")
    # print(f"总帧数: {total_frames}")
    print(f"分辨率: {width}x{height}")
    # print(f"编码格式: {codec_str}")

    # 检查输出文件夹是否存在，不存在则创建
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 文件名，去除路径和后缀
    file_name = os.path.basename(ts_file_path).split('.')[0]

    # 当前帧计数器
    frame_count = 0

    # 需要抽取的帧的序号，间隔为 `interval` 秒
    frame_interval = fps * interval
    count = 0
    while True:
        # 读取视频的帧
        ret, frame = cap.read()

        # 如果读取成功
        if not ret:
            break

        if frame_count % frame_interval == 0:
            # 构建输出图片的路径
            output_path = os.path.join(output_dir, f"{file_name}_{count}.jpg")
            # 保存当前帧为 JPG 图片
            cv2.imwrite(output_path, frame)
            print(f"已保存帧:{frame_count} 到 {output_path} ")
            count += 1
            # 屏幕打印frame
            # cv2.imshow('frame', frame)

        # 增加帧计数
        frame_count += 1

    # 释放视频对象
    cap.release()


def extract_frames_from_all_ts_files(input_dir, output_root_dir, interval=3):
    # 遍历指定目录下的所有文件
    for root, dirs, files in os.walk(input_dir):
        # 遍历每个文件
        for file in files:
            # 如果是 TS 文件
            if file.lower().endswith(".ts"):
                ts_file_path = os.path.join(root, file)

                # 构建输出目录（保持输入目录的结构）
                relative_dir = os.path.relpath(root, input_dir)
                output_dir = os.path.join(output_root_dir, relative_dir)

                # 调用提取帧的函数
                extract_frames(ts_file_path, output_dir, interval)


# 示例用法
input_dir = "/media/chenglei3/77D014CED257D1411/data/车舱数据"  # 需要遍历的TS文件所在目录
output_root_dir = "/media/chenglei3/77D014CED257D1411/data/车舱数据/frames"  # 提取帧后保存的根目录

# 遍历目录并提取所有TS文件的帧
extract_frames_from_all_ts_files(input_dir, output_root_dir)
